import math
import time
import numpy as np
import logging
from XRobot.envs import RobotEnv
from XRobot.commons.cv_utils import get_cam_intrinsic
import XRobot.commons.transform as T

logging.basicConfig(level=logging.INFO)
import cv2


def main():

    robot = 'UR5e'
    # robot = 'Panda'

    env = RobotEnv(
        robot=robot+'Grasp',
        render_mode='human',
        control_freq=200,
        controller='CARTIK',
        is_show_camera_in_cv=True,
        is_render_camera_offscreen=True,
        camera_in_render="0_cam",
        # camera_in_window="frontview",     # 固定的渲染视角
    )

    ## TODO: 示例代码，可以删除，并在此编写你的代码

    start_time = time.time()
    env.reset()

    action_init = np.array([0.3, 0.4, 0.4, 1, 0, 0, 0])
    for t in range(int(20)):
        env.step(action_init)

    cv_image = env.render("rgb_array")
    cv2.imwrite(f"image.png", cv_image)

    q = T.euler_2_quat(np.array([0, 0, 45 * np.pi / 180]))
    action = np.array([0.4, 0.3, 0.1, q[3], q[0], q[1], q[2]])

    # 运动到物块位置
    for t in range(int(200)):
        env.robot.end['arm0'].open()
        env.step(action)

    # 夹爪向下运动
    action[2] = 0.02
    for t in range(int(100)):
        env.step(action)

    # 夹爪闭合
    for t in range(int(20)):
        env.robot.end['arm0'].close()
        env.step(action)

    # 夹爪向上运动
    action[2] = 0.1
    for t in range(int(100)):
        env.step(action)

    # 运动到对应的圆盘位置
    action = np.array([0.4, -0.4, 0.1, q[3], q[0], q[1], q[2]])
    for t in range(int(200)):
        env.step(action)

    # 末端旋转
    action[3:] = np.array([0.707, 0, 0, 0.707])
    for t in range(int(100)):
        env.step(action)

    # 打开夹爪
    for t in range(int(20)):
        env.robot.end['arm0'].open()
        env.step(action)

    # 回到中间位置
    action[1] = 0
    for t in range(int(100)):
        env.step(action)

    task_time = time.time() - start_time
    print("time: %.2f" % task_time)

    env.close()


if __name__ == "__main__":
    main()



